import asyncio
import os

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from pydantic import SecretStr

from browser_use import Agent, Browser
from aioconsole import ainput  # 需要安装 aioconsole

# dotenv
load_dotenv()

api_key = os.getenv('api_key', '')
if not api_key:
	raise ValueError('DEEPSEEK_API_KEY is not set')


async def run_pentest():
    
    browser = Browser()
    llm=ChatOpenAI(
                base_url='https://api.deepseek.com/v1',
                model='deepseek-chat',
                api_key=SecretStr(api_key))
    async with await browser.new_context() as context:
        agent = Agent(
            task=(
                '1. 连接http://192.168.65.97/dvwa/index.php '
                '2. 这是一个DVWA的靶机，使用默认密码登录'
		'3. 进入DVWA Security,将安全级别设为low,然后提交'
                     
                
            ),
            llm=llm,
            
            use_vision=False,
            browser_context=context,
        )

        await agent.run()
        while True:
            user_input = input("请输入任务内容（输入'q'退出）: ")
    
            if user_input.lower() == 'q':
                print("程序已退出")
                break
    
            agent.add_new_task(user_input)
            await agent.run()
        
	

if __name__ == '__main__':
	asyncio.run(run_pentest())
